#include "AnaDef.h"
#include "XGLUtils.h"

void doQs(const Int_t opt)
{
  TTree * tree = (TTree*)gDirectory->Get("newdraw_Tree");
  if(!tree){
    printf("no tree!!\n");
    gDirectory->ls();
    exit(1);
  }

  Double_t ltbg = -999;
  Double_t tpcsig = -999;
  Int_t mag = -999;
  Int_t run = -999;
  Int_t nch = -999;
  Int_t itrk = -999;
  Int_t ifill = -999;
  Int_t spe = -999;

  tree->SetBranchAddress("ltbg", &ltbg);
  tree->SetBranchAddress("tpcsig",&tpcsig);
  tree->SetBranchAddress("mag",&mag);
  tree->SetBranchAddress("run",&run);
  tree->SetBranchAddress("nchamber",&nch);
  tree->SetBranchAddress("itrk",&itrk);
  tree->SetBranchAddress("ifill",&ifill);
  tree->SetBranchAddress("spe",&spe);

  //===========================

  Int_t trdncls = -999, tpcncls = -999, concls = -999;
  TVectorD *vecq=0x0;
  TVectorD *vecx=0x0;

  tree->SetBranchAddress("trdncls0",&trdncls);
  tree->SetBranchAddress("tpcncls0",&tpcncls);
  tree->SetBranchAddress("concls0",&concls);
  tree->SetBranchAddress("coarrayQ0",&vecq);
  tree->SetBranchAddress("coarrayX0",&vecx);

  Double_t tmean = -999;
  if(opt==0){
    tree->SetBranchAddress("tpcQ0",&tmean);
  }
  else if(opt==10){
    tree->SetBranchAddress("trdQ0",&tmean);
  }
  else if(opt==20){
    tree->SetBranchAddress("coQ0",&tmean);
  }
  else{
    printf("opt error %d\n", opt); exit(1);
  }

  Double_t lowFrac = -999, highFrac = -999;
  lowFrac = 0; highFrac = 0.6;

  // tree->SetBranchAddress("",&);
  //============================================================

  const Int_t nx = 50;
  const Double_t xmin = -0.5;
  const Double_t xmax = 4.5;

  Double_t tmax = -999;
  if(opt>=0 && opt<10) tmax = 160;
  else if(opt>=10 && opt<20) tmax = 32;
  else tmax = 346;

  const Double_t tmin = -1*tmax;
  const Int_t nt = tmax-tmin;

  const Int_t nq = 10000;
  const Double_t qmax = 20;

  TH3D *h3raw=new TH3D("h3raw","", nx, xmin, xmax, nt, tmin, tmax, nq, 0, qmax);
  TH3D *h3tru=new TH3D("h3tru","", nx, xmin, xmax, nt, tmin, tmax, nq, 0, qmax);

  TH2D *htbraw = new TH2D("htbraw","", nx, xmin, xmax, nt, tmin, tmax);
  TH2D *htbtru = new TH2D("htbtru","", nx, xmin, xmax, nt, tmin, tmax);

  const Int_t ntot = tree->GetEntries();
  for(Int_t ii=0; ii<ntot; ii++){
    tree->GetEntry(ii);
    if(ii%100000==0){
      printf("Entry: %d/%d\n",ii, ntot);
    }

    if(nch<6) continue;//select only n==6, important to effect of truncation
    if(ifill!=0) continue;
    if(fabs(mag)==1 && ltbg>1.45) continue;
    if(fabs(mag)==5 && ltbg<1.5) continue;
    //if(fabs(mag)==5 && itrk!=ktrk) continue;
    if(fabs(mag)==1 && tpcsig>44 && ltbg<1.5) continue;
    if(fabs(mag)==5 && run<150000 && ((tpcsig>160&&ltbg<1.5)||(tpcsig>59&&tpcsig<90&&ltbg<1.6))) continue;
    if(ltbg<xmin || ltbg>xmax) continue;
    if(spe!=1) continue;

    //============
    TVectorD * arrayq = 0x0;
    TVectorD * arrayx = 0x0;
    Bool_t knew = kFALSE;
    if(opt>=20){
      arrayq= vecq; 
      arrayx = vecx;
    }
    else{
      knew = kTRUE;

      arrayq = new TVectorD(200);
      arrayx = new TVectorD(200);

      if(opt>=0&&opt<10){
        for(Int_t itrd=0; itrd<tpcncls; itrd++){
          (*arrayq)[itrd]=(*vecq)[itrd];
          (*arrayx)[itrd]=(*vecx)[itrd];
        }
      } 
      else if(opt>=10&&opt<20){
        for(Int_t itrd=0; itrd<trdncls; itrd++){
          (*arrayq)[itrd]=(*vecq)[tpcncls+itrd];
          (*arrayx)[itrd]=(Int_t)((*vecx)[tpcncls+itrd]-159)%31;
        }
      }
    }

    //printf("%d %f %f\n", ii, tmean, ltbg);
    Int_t nq=0;
    Double_t tts[1000], qqs[1000];

    for(Int_t ir=0; ir<arrayq->GetNrows(); ir++){
      const Double_t qq = (*arrayq)[ir];
      const Double_t rawt = (*arrayx)[ir];//*6+10;
      Double_t tt = rawt+h3raw->GetYaxis()->GetBinWidth(1);
      if(itrk==0)
        tt*=-1;

      if(qq>1e-12){
        h3raw->Fill(ltbg, tt, qq);
        htbraw->Fill(ltbg,tt);

        tts[nq] = tt;
        qqs[nq] = qq;
        nq++;
      }
    }

    //------

    Int_t idx[1000];
    TMath::Sort(nq, qqs, idx, 0);

    const Int_t istart = lowFrac*nq;
    const Int_t istop = highFrac*nq;

    Double_t sum=0;
    Int_t nused = 0;
    for(Int_t iso=istart; iso<istop; iso++){
      const Int_t ids = idx[iso];

      h3tru->Fill(ltbg, tts[ids], qqs[ids]);        
      htbtru->Fill(ltbg, tts[ids]);
      sum+=qqs[ids];
      nused++;
    }

    Double_t guess = sum/nused;

    if(fabs(guess-tmean)>1e-12){
      printf("wrong summation! %e %e %d -- %e\n", tmean, sum, nused, guess-tmean); exit(1);
    }

    if(knew){
      delete arrayq;
      delete arrayx;
    }
  }

  htbraw->Sumw2();
  htbtru->Sumw2();

  TH2D *htbratio = (TH2D*) htbtru->Clone("htbratio");
  //htbratio->Divide(htbraw);
  for(Int_t ix=0; ix<=htbratio->GetNbinsX()+1; ix++){
    for(Int_t iy=0; iy<=htbratio->GetNbinsY()+1; iy++){
      htbratio->SetBinContent(ix,iy,0);
      htbratio->SetBinError(ix,iy,0);
      const Double_t nraw = htbraw->GetBinContent(ix,iy);
      if(nraw<200)
        continue;

      const Double_t ntbtru = htbtru->GetBinContent(ix,iy);
      const Double_t eff = ntbtru/nraw;
      const Double_t err = sqrt(eff*(1-eff)/nraw);
      htbratio->SetBinContent(ix,iy,eff);
      htbratio->SetBinError(ix,iy,err);
    }
  }
  //-----

  TList *list=new TList;

  TH2D *h3rawmean = new TH2D(Form("%smean", h3raw->GetName()),"", nx, xmin, xmax, nt, tmin, tmax);
  TH2D *h3rawrms  = new TH2D(Form("%srms",  h3raw->GetName()),"", nx, xmin, xmax, nt, tmin, tmax);

  TH2D *h3trumean = new TH2D(Form("%smean", h3tru->GetName()),"", nx, xmin, xmax, nt, tmin, tmax);
  TH2D *h3trurms  = new TH2D(Form("%srms",  h3tru->GetName()),"", nx, xmin, xmax, nt, tmin, tmax);

  const Double_t thres = 500;
  for(Int_t ix=1; ix<=nx; ix++){
    printf("doing ix: %d/%d\n", ix, nx);

    for(Int_t it=1; it<=nt; it++){
      TH1D * hzraw = h3raw->ProjectionZ(Form("rawx%dy%d", ix, it), ix,ix,it,it);

      if(hzraw->GetEntries()<thres){
        delete hzraw;
        continue;
      }

      //list->Add(hzraw);
      Double_t rms=-999, err=-999;
      const Double_t mean = XGLUtils::TruncatedMean(hzraw, lowFrac, highFrac, &rms, &err);
      h3rawmean->SetBinContent(ix,it,mean);
      h3rawmean->SetBinError(ix,it,err);

      h3rawrms->SetBinContent(ix,it,rms);
      h3rawrms->SetBinError(ix,it,0);

      //hzraw->Rebin(100);
      delete hzraw;
    }
  }

  for(Int_t ix=1; ix<=nx; ix++){
    printf("doing ix: %d/%d\n", ix, nx);

    for(Int_t it=1; it<=nt; it++){
      TH1D * hztru = h3tru->ProjectionZ(Form("trux%dy%d", ix, it), ix,ix,it,it);

      if(hztru->GetEntries()<thres){
        delete hztru;
        continue;
      }

      //list->Add(hztru);
      Double_t rms=-999, err=-999;
      const Double_t mean = XGLUtils::TruncatedMean(hztru, lowFrac, highFrac, &rms, &err);
      h3trumean->SetBinContent(ix,it,mean);
      h3trumean->SetBinError(ix,it,err);

      h3trurms->SetBinContent(ix,it,rms);
      h3trurms->SetBinError(ix,it,0);

      //hztru->Rebin(100);
      delete hztru;
    }
  }

  const Int_t prof0=h3rawmean->GetNbinsY()/2+1;
  const Int_t prof1=h3rawmean->GetNbinsY();
  TProfile *htbrawmean0 = h3rawmean->ProfileX(Form("htbrawmean0_%d_%d", 1,     prof0), 1,     prof0);
  TProfile *htbrawmean1 = h3rawmean->ProfileX(Form("htbrawmean1_%d_%d", prof0, prof1), prof0, prof1);
  htbrawmean0->Scale(-1);

  TProfile *htbtrumean0 = h3rawmean->ProfileX(Form("htbtrumean0_%d_%d", 1,     prof0), 1,     prof0);
  TProfile *htbtrumean1 = h3rawmean->ProfileX(Form("htbtrumean1_%d_%d", prof0, prof1), prof0, prof1);
  htbtrumean0->Scale(-1);

  TFile *fout=new TFile(Form("vecqout%d.root", opt),"recreate");
  htbraw->Write();
  htbtru->Write();
  htbratio->Write();

  h3rawmean->Write();
  htbrawmean0->Write();
  htbrawmean1->Write();
  h3rawrms->Write();

  h3trumean->Write();
  htbtrumean0->Write();
  htbtrumean1->Write();
  h3trurms->Write();

  list->Write("list", TList::kSingleKey);
  fout->Save();
  fout->Close();
}

int main(int argc, char *argv[])
{
  for(int ii=0; ii<argc; ii++){
    printf("%d: %s\n", ii, argv[ii]);
  }

  if(argc!=3){
    printf("argc!=3\n");
    exit(1);
  }

  TFile::Open(argv[1]);

  doQs(atoi(argv[2]));

  return 0;
}
